import numpy as np
import time
from vrep_api import vrep
from algorithm.mathlib import utilities as mathutl
import multiprocessing as mp



class read_data_package():
    def __init__(self):
        self.forceVector = np.zeros(3)
        self.torqueVector = np.zeros(3)
        self.tip_position = np.zeros(3)
        self.tip_orientation = np.zeros(3)

class send_data_package():
    def __init__(self):
        self.target_position = np.zeros(3)
        self.target_orientation = np.zeros(3)

class sim():
    def __init__(self):

        self.sim_client = None
        self.sim_tip_handle = None
        self.sim_target_handle = None
        self.sim_ft_handle = None
        self.sim_tool_frame_handle = None
        self.sim_sensor_frame_handle = None
        self.read_data = read_data_package()
        self.send_data = send_data_package()

        self.sim_start()
        self.sim_config()

    '''
        This part is for vrep connection.
    '''

    def sim_start(self):
        while True:
            vrep.simxFinish(-1) # Just in case, close all opened connections
            print('[ENVIRONMENT STATE]: Waiting for V-REP ...')
            self.sim_client = vrep.simxStart('127.0.0.1', 19997, True, True, 5000, 5) # be careful of the connectPort

            if self.sim_client > -1:
                print('[ENVIRONMENT STATE]: Connected to simulation.')
                break
            else:
                time.sleep(0.2)
                print('[ENVIRONMENT STATE]: Failed to connect to V-REP. Exiting.')

        time.sleep(2)
    def sim_stop(self):
        vrep.simxStopSimulation(self.sim_client, vrep.simx_opmode_blocking)

    def sim_config(self):
        sim_ret, self.sim_tip_handle = vrep.simxGetObjectHandle(self.sim_client, 'UR5_tip', vrep.simx_opmode_blocking)
        sim_ret, self.sim_target_handle = vrep.simxGetObjectHandle(self.sim_client, 'UR5_target', vrep.simx_opmode_blocking)
        sim_ret, self.sim_ft_handle = vrep.simxGetObjectHandle(self.sim_client, 'UR5_connection', vrep.simx_opmode_blocking)
        sim_ret, self.sim_sensor_frame_handle = vrep.simxGetObjectHandle(self.sim_client, 'sensor_frame', vrep.simx_opmode_blocking)
        sim_ret, self.sim_tool_frame_handle = vrep.simxGetObjectHandle(self.sim_client, 'tool_frame', vrep.simx_opmode_blocking)

    def sim_read_init_data(self):
        sim_ret, pos_tip2sensor = vrep.simxGetObjectPosition(self.sim_client, self.sim_tip_handle, self.sim_sensor_frame_handle, vrep.simx_opmode_blocking)

        sim_ret, ori_tip2sensor = vrep.simxGetObjectOrientation(self.sim_client, self.sim_tip_handle, self.sim_sensor_frame_handle, vrep.simx_opmode_blocking)
        sim_ret, pos_toolcom2sensor = vrep.simxGetObjectPosition(self.sim_client, self.sim_tool_frame_handle, self.sim_sensor_frame_handle, vrep.simx_opmode_blocking)
        print("pos_tip2sensor", pos_tip2sensor)
        print("ori_tip2sensor", ori_tip2sensor)
        print("pos_toolcom2sensor", pos_toolcom2sensor)

    def sim_read_data(self):
        sim_ret, state, forceVector, torqueVector = vrep.simxReadForceSensor(self.sim_client,self.sim_ft_handle,vrep.simx_opmode_streaming)
        sim_ret, UR5_tip_position = vrep.simxGetObjectPosition(self.sim_client, self.sim_tip_handle, -1,
                                                                vrep.simx_opmode_blocking)
        sim_ret, UR5_tip_orientation = vrep.simxGetObjectOrientation(self.sim_client, self.sim_tip_handle, -1,
                                                                vrep.simx_opmode_blocking)
        self.read_data.forceVector = forceVector
        self.read_data.torqueVector = torqueVector
        self.read_data.tip_position = UR5_tip_position
        self.read_data.tip_orientation = UR5_tip_orientation

    def sim_send_data(self):
        vrep.simxSetObjectPosition(self.sim_client, self.sim_target_handle, -1, self.send_data.target_position, vrep.simx_opmode_oneshot)
        vrep.simxSetObjectOrientation(self.sim_client, self.sim_target_handle, -1, self.send_data.target_orientation, vrep.simx_opmode_oneshot)

    def task(self):
        self.sim_read_data()